# _*_ encoding:utf-8 _*_

import matplotlib.pyplot as plt
import random
import optparse
import sys
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


def load_edges(filename):
    tree = ET.ElementTree(file=filename)
    edges = dict()
    for child in tree.getroot():
        if child.tag == 'edge':
            length = child[0].attrib['length']
            id = child.attrib['id']
            edges.setdefault(id, {})
            edges[id].setdefault('length', float(length))
            points = []
            for point in child[0].attrib['shape'].strip().split(" "):
                x, y = map(float, point.split(","))
                points.append([x, y])
            edges[id].setdefault('shape', points)
            edges[id].setdefault('speed', float(child[0].attrib['speed']))
    return edges


def test_load_edges():
    test_filename = "../../data/test_net.net.xml"
    print load_edges(test_filename)


def process_routes(filename, edges):
    tree = ET.ElementTree(file=filename)
    trips = []
    for child in tree.getroot():
        if child.tag == 'vehicle':
            trip = dict()
            trip.setdefault('id', child.attrib['id'])
            # trip.setdefault('type', child.attrib['type'])
            accumulate = 0
            routes = child[0].attrib['edges'].split(" ")
            for edge in routes:
                accumulate += float(edges[edge]["length"])
            trip.setdefault('distance', accumulate)
            trip.setdefault("start_edge", routes[0])
            trip.setdefault("end_edge", routes[-1])
            trips.append(trip)
    return trips


def test_process_routes():
    test_edges_file = "../../data/test_net.net.xml"
    test_routes_file = "../../data/test_route.rou.xml"
    edges = load_edges(test_edges_file)
    print process_routes(test_routes_file, edges)


def plot_distance_distribution(trips, unit):
    distance = {}
    unit *= 1000
    for trip in trips:
        d = int(trip['distance']/unit)
        distance.setdefault(d, 0)
        distance[d] += 1
    sort = sorted(distance.iteritems(), key=lambda z: z[0], reverse=False)
    X = []
    y = []
    for key, value in sort:
        X.append(key)
        y.append(value)
    plt.bar(X, y)
    plt.xlabel('{0}m'.format(unit))
    plt.ylabel('number')
    plt.show()


def plot_map(edges):
    fig = plt.figure()
    ax = fig.add_axes([0.1, 0.1, 0.7, 0.8])
    for key, value in edges.items():
        X = []
        Y = []
        for point in value['shape']:
            X.append(point[0])
            Y.append(point[1])
        ax.plot(X, Y, color='black')
    return fig


def start_point_distribution(edges, trips):
    fig = plot_map(edges)
    X = []
    Y = []
    for trip in trips:
        start_edge = edges[trip['start_edge']]
        x, y = start_edge['shape'][0]
        x += random.gauss(0, 5.0)
        y += random.gauss(0, 5.0)
        # start_points.append((x, y))
        X.append(x)
        Y.append(y)
    ax = fig.axes[0]
    ax.scatter(X, Y, color='red')
    plt.show(fig)


def destination_point_distribution(edges, trips):
    fig = plot_map(edges)
    X = []
    Y = []
    for trip in trips:
        end_edge = edges[trip['end_edge']]
        x, y = end_edge['shape'][-1]
        x += random.gauss(0, 5.0)
        y += random.gauss(0, 5.0)

        X.append(x)
        Y.append(y)
    ax = fig.axes[0]
    ax.scatter(X, Y, color='blue')
    plt.show(fig)


def start_destination_line(edges, trips, n=-1):
    fig = plot_map(edges)
    ax = fig.axes[0]
    lines = []
    if n == -1:
        n = len(trips)
    for i, trip in enumerate(trips):
        start_edge = edges[trip['start_edge']]
        end_edge = edges[trip['end_edge']]
        sx, sy = start_edge['shape'][0]
        ex, ey = end_edge['shape'][-1]
        sx += random.gauss(0, 2.0)
        sy += random.gauss(0, 2.0)
        ex += random.gauss(0, 2.0)
        sy += random.gauss(0, 2.0)
        if len(lines) < n:
            lines.append(([sx, ex], [sy, ey]))
        else:
            k = random.randint(0, i - 1)
            if k < n:
                lines[k] = ([sx, ex], [sy, ey])
    for line in lines:
        ax.plot(line[0], line[1])
    plt.show(fig)


def plot_traffic_status(edges):
    fig = plot_map(edges)
    ax = fig.axes[0]
    for key, value in edges.items():
        X = []
        Y = []
        for point in value['shape']:
            X.append(point[0])
            Y.append(point[1])
        prob = random.random()
        if prob < 0.2:
            ax.plot(X, Y, color='#458b00')
        elif prob < 0.4:
            ax.plot(X, Y, color='#66cd00')
        elif prob < 0.6:
            ax.plot(X, Y, color="#eeee00")
        elif prob < 0.8:
            ax.plot(X, Y, color="#ee7600")
        else:
            ax.plot(X, Y, color="#ee2c2c")
    plt.show(fig)


def get_options(args=None):
    optParser = optparse.OptionParser()
    optParser.add_option("-n", "--net-file", dest='net_file', help="the define net file")
    optParser.add_option("-r", "--rout-file", dest="route_file", help="the route file to analyse")
    optParser.add_option("-d", "--distance", type="float", dest="distance", help="plot the distance distribution based on\
     the given value, the unit is KM")
    optParser.add_option("-s", "--start", type="int", dest="start", help="Plot the start point distribution")
    optParser.add_option("-e", "--end", type="int", dest="end", help="Plot the destination point distribution \
                                                    start and end should not be set the same time")
    optParser.add_option("-l", "--line", type='int', dest='line', help="plot the od line in the map. if p is given \
                        then sample p lines from the od and plot the p lines")
    optParser.add_option("-p", "--pnumber", type="int", dest='number', help="Plot number lines")
    (options, args) = optParser.parse_args(args=args)
    if not options.net_file:
        optParser.print_help()
        sys.exit()

    if not options.route_file:
        optParser.print_help()
        sys.exit()

    if options.distance and options.distance <= 0:
        print "option distance should greater than 0"

    if not (options.start is None) and not (options.end is None):
        print "--start and --end should not be set at the same time"

    return options


def main(options):
    # net_file = "../../data/hangzhou.net.xml"
    # routes_file = "../../data/10221515.rou.xml"
    edges = load_edges(options.net_file)
    trips = process_routes(options.route_file, edges)
    print len(edges)
    print len(trips)
    if not (options.distance is None ):
        plot_distance_distribution(trips, options.distance)
    start = 0
    if not (options.start is None):
        start = options.start
    if start:
        start_point_distribution(edges, trips)

    end = 0
    if not (options.end is None):
        end = options.end
    if end:
        destination_point_distribution(edges, trips)

    line = 0
    number = -1
    if options.number is not None:
        number = options.number
    if options.line is not None and options.line:
        start_destination_line(edges, trips, number)
    # statistics the trip distance
    # plot_distance_distribution(trips)
    # fig = plot_map(edges)
    # plt.show(fig)
    # start_point_distribution(edges, trips)
    # destination_point_distribution(edges, trips)
    # start_destination_line(edges, trips, 100)
    # plot_traffic_status(edges)

if __name__ == '__main__':
    main(get_options())
    # test_process_routes()
    # test_load_edges()
